package com.example.demo1122.utils;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.Constants;
import com.alibaba.dashscope.utils.JsonUtils;
import com.example.demo1122.entity.Question;
import io.reactivex.Flowable;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class QwenStreamCall {
    public static String streamCallWithMessage(String sysRole, List<Question> questionList, String questionContent)
            throws NoApiKeyException, ApiException, InputRequiredException {
        //TODO：优化这一块的keyList，现在每次调用都要生成一次
        ArrayList<String> keyList = new ArrayList<>();
        keyList.add("sk-d13928106aaf467788082be947c326ad");
        keyList.add("sk-d18f90ccfa964b9c8c407b32b831cc4d");
        keyList.add("sk-a0d5ce73f0884ad092a6b14295d6106c");
        Integer random = new Random().nextInt(3);
        Constants.apiKey = keyList.get(random);
        Generation gen = new Generation();
        // append历史question answer后续添加
        ArrayList<Message> msgList = new ArrayList<>();
        msgList.add(Message
                .builder()
                .role(Role.SYSTEM.getValue())
                .content(sysRole)
                .build());
        for (Question question_answerPair : questionList) {
            msgList.add(Message
                    .builder()
                    .role(Role.USER.getValue())
                    .content(question_answerPair.getQuestionContent())
                    .build());
            msgList.add(Message
                    .builder()
                    .role(Role.ASSISTANT.getValue())
                    .content(question_answerPair.getAnswerContent())
                    .build());
        }
        msgList.add(Message
                .builder()
                .role(Role.USER.getValue())
                .content(questionContent)
                .build());
        QwenParam param =
                QwenParam.builder()
                        .model(Generation.Models.QWEN_PLUS)
                        .messages(msgList)
                        .resultFormat(QwenParam.ResultFormat.MESSAGE)
                        .topP(0.8)
                        .enableSearch(true)
                        .incrementalOutput(true) // get streaming output incrementally
                        .build();
        Flowable<GenerationResult> result = gen.streamCall(param);
        StringBuilder fullContent = new StringBuilder();
        // blockingForEach,在每一次响应后block获取输出append
        result.blockingForEach(message -> {
            fullContent.append(message.getOutput().getChoices().get(0).getMessage().getContent());
            System.out.println(JsonUtils.toJson(message));
        });
        System.out.println("Full content: \n" + fullContent.toString());
        return fullContent.toString();
    }

    public static Flowable<GenerationResult> streamCallWithMessage_stream(String sysRole, List<Question> questionList, String questionContent)
            throws NoApiKeyException, ApiException, InputRequiredException {
        Constants.apiKey = "sk-3e7738dd94e34b7a8d9d4747e6730afd"; // 后续将api key放入配置文件中
        Generation gen = new Generation();
        // 其中关于assistant设置需要后续添加
        // append历史question answer后续添加
        ArrayList<Message> msgList = new ArrayList<>();
        msgList.add(Message
                .builder()
                .role(Role.SYSTEM.getValue())
                .content(sysRole)
                .build());
        for (Question question_answerPair : questionList) {
            msgList.add(Message
                    .builder()
                    .role(Role.USER.getValue())
                    .content(question_answerPair.getQuestionContent())
                    .build());
            msgList.add(Message
                    .builder()
                    .role(Role.ASSISTANT.getValue())
                    .content(question_answerPair.getAnswerContent())
                    .build());
        }
        msgList.add(Message
                .builder()
                .role(Role.USER.getValue())
                .content(questionContent)
                .build());
        QwenParam param =
                QwenParam.builder()
                        .model(Generation.Models.QWEN_PLUS)
                        .messages(msgList)
                        .resultFormat(QwenParam.ResultFormat.MESSAGE)
                        .topP(0.8)
                        .enableSearch(true)
                        .incrementalOutput(true) // get streaming output incrementally
                        .build();
        Flowable<GenerationResult> result = gen.streamCall(param);
        return result;
    }
}


